#!/usr/bin/env python3
"""
Tushare数据源实现
提供免费的A股数据获取功能
"""

import pandas as pd
from datetime import datetime, timedelta
from typing import Dict, Any, Optional
import time

from .base import BaseDataSource, DataSourceStatus
from ..utils.logger import get_logger

logger = get_logger("TushareSource")


class TushareSource(BaseDataSource):
    """Tushare数据源"""

    def __init__(self, config: Dict[str, Any] = None):
        super().__init__("tushare", config)
        self.ts = None
        self.pro = None
        self._initialize_fetcher()

    def _initialize_fetcher(self):
        """初始化数据获取器"""
        try:
            import tushare as ts
            self.ts = ts
            
            # 尝试使用免费接口
            token = self.config.get('token', '')
            if token:
                ts.set_token(token)
                self.pro = ts.pro_api()
                logger.info("Tushare Pro API初始化成功")
            else:
                logger.info("使用Tushare免费接口")
            
            self.update_status(DataSourceStatus.AVAILABLE)
            logger.info("Tushare数据源初始化成功")
        except ImportError:
            self.update_status(DataSourceStatus.ERROR, "tushare未安装")
            logger.error("tushare未安装，请运行: pip install tushare")
        except Exception as e:
            self.update_status(DataSourceStatus.ERROR, str(e))
            logger.error(f"Tushare数据源初始化失败: {e}")

    def check_availability(self) -> bool:
        """检查数据源是否可用"""
        try:
            if not self.ts:
                self.update_status(DataSourceStatus.ERROR, "tushare未初始化")
                return False

            # 简单检查，不进行网络请求
            available = self.ts is not None
            self.update_status(
                DataSourceStatus.AVAILABLE if available else DataSourceStatus.ERROR,
                None if available else "tushare未可用"
            )
            return available
        except Exception as e:
            self.update_status(DataSourceStatus.ERROR, str(e))
            return False

    def get_stock_list(self) -> pd.DataFrame:
        """获取股票列表"""
        try:
            if self.pro:
                # 使用Pro接口
                df = self.pro.stock_basic(exchange='', list_status='L', 
                                        fields='ts_code,symbol,name,area,industry,market')
                if not df.empty:
                    # 转换为统一格式
                    df['symbol'] = df['symbol'].astype(str).str.zfill(6)
                    df = df.rename(columns={'ts_code': 'code'})
                    df['data_source'] = 'tushare'
                return df
            else:
                # 使用免费接口
                df = self.ts.get_stock_basics()
                if not df.empty:
                    df = df.reset_index()
                    df = df.rename(columns={'code': 'symbol'})
                    df['symbol'] = df['symbol'].astype(str).str.zfill(6)
                    df['data_source'] = 'tushare'
                return df
        except Exception as e:
            logger.error(f"Tushare获取股票列表失败: {e}")
            return pd.DataFrame()

    def fetch_history_data(self, symbol: str, interval: str,
                          start_date: str, end_date: str) -> pd.DataFrame:
        """获取历史数据"""
        try:
            # 转换股票代码格式
            if len(symbol) == 6 and symbol.isdigit():
                if symbol.startswith('6'):
                    ts_code = f"{symbol}.SH"
                else:
                    ts_code = f"{symbol}.SZ"
            else:
                ts_code = symbol

            if self.pro:
                # 使用Pro接口
                df = self.pro.daily(ts_code=ts_code, start_date=start_date.replace('-', ''),
                                  end_date=end_date.replace('-', ''))
                if not df.empty:
                    # 转换为统一格式
                    df['Date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d')
                    df = df.rename(columns={
                        'open': 'Open', 'high': 'High', 'low': 'Low', 
                        'close': 'Close', 'vol': 'Volume'
                    })
                    df = df[['Date', 'Open', 'High', 'Low', 'Close', 'Volume']].sort_values('Date')
            else:
                # 使用免费接口
                df = self.ts.get_hist_data(symbol, start=start_date, end=end_date)
                if not df.empty:
                    df = df.reset_index()
                    df['Date'] = pd.to_datetime(df['date'])
                    df = df.rename(columns={
                        'open': 'Open', 'high': 'High', 'low': 'Low', 
                        'close': 'Close', 'volume': 'Volume'
                    })
                    df = df[['Date', 'Open', 'High', 'Low', 'Close', 'Volume']].sort_values('Date')

            # 添加延迟避免频率限制
            time.sleep(0.1)
            return df

        except Exception as e:
            logger.error(f"Tushare获取历史数据失败 {symbol}: {e}")
            return pd.DataFrame()

    def get_stock_info(self, symbol: str) -> Dict[str, Any]:
        """获取股票信息"""
        try:
            if self.pro:
                # 转换股票代码格式
                if len(symbol) == 6 and symbol.isdigit():
                    if symbol.startswith('6'):
                        ts_code = f"{symbol}.SH"
                    else:
                        ts_code = f"{symbol}.SZ"
                else:
                    ts_code = symbol

                df = self.pro.stock_basic(ts_code=ts_code, fields='ts_code,symbol,name,area,industry,market')
                if not df.empty:
                    info = df.iloc[0].to_dict()
                    return {
                        'symbol': symbol,
                        'name': info.get('name', 'Unknown'),
                        'industry': info.get('industry', 'Unknown'),
                        'area': info.get('area', 'Unknown'),
                        'market': info.get('market', 'Unknown')
                    }
            else:
                # 使用免费接口
                basics = self.ts.get_stock_basics()
                if symbol in basics.index:
                    info = basics.loc[symbol]
                    return {
                        'symbol': symbol,
                        'name': info.get('name', 'Unknown'),
                        'industry': info.get('industry', 'Unknown'),
                        'area': info.get('area', 'Unknown')
                    }
            return {}
        except Exception as e:
            logger.error(f"Tushare获取股票信息失败 {symbol}: {e}")
            return {}
